# System Import
from enum import Enum
from importlib.machinery import WindowsRegistryFinder

# ThirdParty Import
import taichi as Ti
import numpy as np
from Kernel import Kernel

Ti.init(arch=Ti.gpu)
# Local Import
import Math
from Solver import *
import Config.GlobalParameters as GP

# Ti.init(arch=Ti.gpu, kernel_profiler=True)

VelocityBuffer1 = Ti.Vector.field(2, float, shape=(GP.FluidsResolution, GP.FluidsResolution))
VelocityBuffer2 = Ti.Vector.field(2, float, shape=(GP.FluidsResolution, GP.FluidsResolution))
# Scalar Field
PressureBuffer1 = Ti.field(float, shape=(GP.FluidsResolution, GP.FluidsResolution))
PressureBuffer2 = Ti.field(float, shape=(GP.FluidsResolution, GP.FluidsResolution))

DensityBuffer1 = Ti.field(float, shape=(GP.FluidsResolution, GP.FluidsResolution))
DensityBuffer2 = Ti.field(float, shape=(GP.FluidsResolution, GP.FluidsResolution))
TemperatureBuffer1 = Ti.field(float, shape=(GP.FluidsResolution, GP.FluidsResolution))
TemperatureBuffer2 = Ti.field(float, shape=(GP.FluidsResolution, GP.FluidsResolution))

VelocityCurl = Ti.field(float, shape=(GP.FluidsResolution, GP.FluidsResolution))
VelocityDivBuffer = Ti.field(float, shape=(GP.FluidsResolution, GP.FluidsResolution))
RenderingBuffer = Ti.Vector.field(3, float, shape=(GP.FluidsResolution, GP.FluidsResolution))

MyKernel = Kernel.Kernel(
    kernal_path='./Kernel/KernelDataBase',
    max_rank=GP.MaxRank,
    jacobi_iters=GP.JacobiIters,
    kappa_temperature=GP.TemperatureDiffuseFactor,
    kappa_density=GP.DensityDiffuseFactor,
    kappa_viscosity=GP.ViscosityFactor,
    delta_time=GP.DeltaTime,
    delta_scale=GP.FluidsScale,
    debug_mode=False
)


class TexPair:
    def __init__(self, cur, nxt):
        self.cur = cur
        self.nxt = nxt

    def swap(self):
        self.cur, self.nxt = self.nxt, self.cur


VelocityBuffer = TexPair(VelocityBuffer1, VelocityBuffer2)
PressureBuffer = TexPair(PressureBuffer1, PressureBuffer2)
DensityBuffer = TexPair(DensityBuffer1, DensityBuffer2)
TemperatureBuffer = TexPair(TemperatureBuffer1, TemperatureBuffer2)


# Global Class
class MouseDataGenerator:
    def __init__(self):
        self.prev_mouse = None
        self.prev_color = None

    def __call__(self, gui):
        # [0:2]: normalized delta direction
        # [2:4]: current mouse xy
        # [4:7]: color
        mouse_data = np.zeros(8, dtype=np.float32)
        if gui.is_pressed(Ti.GUI.LMB):
            mxy = np.array(gui.get_cursor_pos(), dtype=np.float32) * GP.WindowResolution
            if self.prev_mouse is None:
                self.prev_mouse = mxy
                # Set lower bound to 0.3 to prevent too dark colors
                self.prev_color = (np.random.rand(3) * 0.7) + 0.3
            else:
                mdir = mxy - self.prev_mouse
                mdir = mdir / (np.linalg.norm(mdir) + 1e-5)
                mouse_data[0], mouse_data[1] = mdir[0], mdir[1]
                mouse_data[2], mouse_data[3] = mxy[0], mxy[1]
                mouse_data[4:7] = self.prev_color
                self.prev_mouse = mxy
        else:
            self.prev_mouse = None
            self.prev_color = None
        return mouse_data


class SolverType(Enum):
    Jacobi = 0,
    RBGS = 1,


'''
https://baike.baidu.com/item/%E9%BE%99%E6%A0%BC%E5%BA%93%E5%A1%94%E6%B3%95/3016350?fr=aladdin
本质上是解其时间导数只与自身与时间有关的偏微分方程，这样就可以分段求梯度拟合Path
'''


# Runge-Kutta method
@Ti.func
def BackTrace(VelocityBuffer: Ti.template(), Position, DeltaTime: Ti.template()):
    # Velocity 是梯度函数记作K， P 是变量函数
    # 按理来说应该是
    # 1. Position = Position * GP.FluidsScale -> 世界坐标
    # 2. Position = Position - 0.5 * DeltaTime * Velocity -> 世界坐标变化
    # 3. Position = Position / GP.FluidsScale -> 局部坐标
    TempPosition = Position
    K1 = LerpFunction(VelocityBuffer, TempPosition)
    TempPosition = Position - 0.5 * DeltaTime * K1 * GP.InvDeltaX
    K2 = LerpFunction(VelocityBuffer, TempPosition)
    TempPosition = Position - 0.5 * DeltaTime * K2 * GP.InvDeltaX
    K3 = LerpFunction(VelocityBuffer, TempPosition)
    TempPosition = Position - 1.0 * DeltaTime * K3 * GP.InvDeltaX
    K4 = LerpFunction(VelocityBuffer, TempPosition)
    TempPosition = Position - DeltaTime * (K1 + 2 * K2 + 2 * K3 + K4) / 6 * GP.InvDeltaX
    return TempPosition


@Ti.func
def CalcArtificialCurlForce(VelocityCurl, i, j):
    L = Math.Sample(VelocityCurl, i - 1, j)
    R = Math.Sample(VelocityCurl, i + 1, j)
    B = Math.Sample(VelocityCurl, i, j - 1)
    T = Math.Sample(VelocityCurl, i, j + 1)
    C = Math.Sample(VelocityCurl, i, j)
    Force = Ti.Vector([abs(T) - abs(B),
                       abs(L) - abs(R)]).normalized(1e-3)
    Force *= GP.AritificalCurlStrength * C
    return Force


# Taichi Kernel
@Ti.kernel
def AddForce(VelocityBuffer: Ti.template(), TemperatureBuffer: Ti.template(), DensityBuffer: Ti.template(),
             VelocityCurl: Ti.template()):
    GravityDir = Ti.Vector([0, -1])
    for i, j in VelocityBuffer:
        Gravity = GP.GravityRatio * DensityBuffer[i, j] * GravityDir * 9.8
        Bouyancy = GP.BouyancyRatio * (TemperatureBuffer[i, j] - GP.AmbientTemperature) * (-GravityDir)
        CurlForce = CalcArtificialCurlForce(VelocityCurl, i, j)

        Momentum = (Gravity + Bouyancy + CurlForce) * GP.DeltaTime

        Velocity = VelocityBuffer[i, j]
        VelocityBuffer[i, j] = Velocity + Momentum


# Emitter Parameters (World Coord!!!)
# Emitter的Check，可以空间换时间，不过这里暂时用if吧，也挺快	
EmitterRadius = 10
EmitterPosition = Ti.Vector([GP.WindowResolution / 2, GP.WindowResolution / 2])
EmitterNoneUniformScale = Ti.Vector([2, 1])
EmitterDensity = 1
EmitterTemperature = 500

LerpFunction = Math.Cubiclerp


@Ti.kernel
def AddEmitter(VelocityBuffer: Ti.template(), TemperatureBuffer: Ti.template(), DensityBuffer: Ti.template()):
    # World to Local
    EmitterRadiusLocal = EmitterRadius / GP.FluidsScale
    EmitterPositionLocal = EmitterPosition / GP.FluidsScale

    for i, j in VelocityBuffer:
        Delta = EmitterPositionLocal - Ti.Vector([i, j])
        Distance = \
            Delta[0] * Delta[0] / (EmitterNoneUniformScale[0] * EmitterNoneUniformScale[0]) + \
            Delta[1] * Delta[1] / (EmitterNoneUniformScale[1] * EmitterNoneUniformScale[1])

        if Distance <= EmitterRadiusLocal * EmitterRadiusLocal:
            DensityBuffer[i, j] = EmitterDensity
            TemperatureBuffer[i, j] = EmitterTemperature


# LerpFunction必须，不然效果依像素分离
@Ti.kernel
def Advect(VelocityBuffer: Ti.template(), AdvectBuffer: Ti.template(), AdvectResultBuffer: Ti.template()):
    for i, j in VelocityBuffer:
        Position = Ti.Vector([i, j]) + 0.5
        Position = BackTrace(VelocityBuffer, Position, GP.DeltaTime)
        # AdvectResultBuffer[i, j] = Math.Sample(AdvectBuffer, Position[0], Position[1])
        AdvectResultBuffer[i, j] = LerpFunction(AdvectBuffer, Position)  # * dye_decay


@Ti.kernel
def Projection(VelocityBuffer: Ti.template(), PressureBuffer: Ti.template()):
    for i, j in VelocityBuffer:
        pl = Math.Sample(PressureBuffer, i - 1, j)
        pr = Math.Sample(PressureBuffer, i + 1, j)
        pb = Math.Sample(PressureBuffer, i, j - 1)
        pt = Math.Sample(PressureBuffer, i, j + 1)
        VelocityBuffer[i, j] -= 0.5 * Ti.Vector([pr - pl, pt - pb]) * GP.InvDeltaX


@Ti.kernel
def Render(RenderingBuffer: Ti.template(), TemperatureBuffer: Ti.template(), DensityBuffer: Ti.template()):
    for i, j in RenderingBuffer:
        RenderingBuffer[i, j] = Ti.Vector([1 - DensityBuffer[i, j], 1 - DensityBuffer[i, j], 1 - DensityBuffer[i, j]]) \
                                + Ti.Vector([
            Ti.max(0, TemperatureBuffer[i, j] - GP.AmbientTemperature - 50) / 100,
            Ti.max(0, TemperatureBuffer[i, j] - GP.AmbientTemperature - 50) / 200,
            Ti.max(0, TemperatureBuffer[i, j] - GP.AmbientTemperature - 50) / 200])


@Ti.kernel
def AddArtificialCurl(VelocityBuffer: Ti.template(), VelocityCurl: Ti.template()):
    for i, j in VelocityBuffer:
        L = Math.Sample(VelocityBuffer, i - 1, j)
        R = Math.Sample(VelocityBuffer, i + 1, j)
        B = Math.Sample(VelocityBuffer, i, j - 1)
        T = Math.Sample(VelocityBuffer, i, j + 1)
        VelocityCurl[i, j] = (R.y - L.y - T.x + B.x) * 0.5


# Function
def SystemStep(SolverMode):
    # w0 -> w1
    AddArtificialCurl(VelocityBuffer.cur, VelocityCurl)
    AddForce(VelocityBuffer.cur, TemperatureBuffer.cur, DensityBuffer.cur, VelocityCurl)
    # Emitter
    AddEmitter(VelocityBuffer.cur, TemperatureBuffer.cur, DensityBuffer.cur)
    # w1 -> w2
    Advect(VelocityBuffer.cur, VelocityBuffer.cur, VelocityBuffer.nxt)
    Advect(VelocityBuffer.cur, DensityBuffer.cur, DensityBuffer.nxt)
    Advect(VelocityBuffer.cur, TemperatureBuffer.cur, TemperatureBuffer.nxt)
    VelocityBuffer.swap()
    DensityBuffer.swap()
    TemperatureBuffer.swap()
    # w2 -> w3
    if SolverMode == SolverType.Jacobi:
        if GP.CompactPoisson == False:
            JacobiSolver.GeneralJacobiSolver(JacobiSolver.Diffuse, VelocityBuffer)
            JacobiSolver.GeneralJacobiSolver(JacobiSolver.ScalarDiffuse, DensityBuffer, MyKernel.DensityDiffuseFactor)
            JacobiSolver.GeneralJacobiSolver(JacobiSolver.ScalarDiffuse, TemperatureBuffer, MyKernel.TemperatureDiffuseFactor)
        else:
            JacobiSolver.GeneralJacobiSolver(
                JacobiSolver.CompactPoisson, None,
                DensityBuffer.cur, DensityBuffer.cur, DensityBuffer.nxt,
                JacobiSolver.TempScalarBuffer1, JacobiSolver.TempScalarBuffer2,
                MyKernel.density_diffuse_weight_2D, MyKernel.density_diffuse_vector_2D,
                None, None)
            JacobiSolver.GeneralJacobiSolver(
                JacobiSolver.CompactPoisson, None,
                VelocityBuffer.cur, VelocityBuffer.cur, VelocityBuffer.nxt,
                JacobiSolver.TempVectorBuffer1, JacobiSolver.TempVectorBuffer2,
                MyKernel.velocity_diffuse_weight_2D, MyKernel.velocity_diffuse_vector_2D,
                None, None)
            # GP.VelocityInitialDiffuseWeights2D, GP.VelocityInitialDiffuseVectors2D)
            JacobiSolver.GeneralJacobiSolver(
                JacobiSolver.CompactPoisson, None,
                TemperatureBuffer.cur, TemperatureBuffer.cur, TemperatureBuffer.nxt,
                JacobiSolver.TempScalarBuffer1, JacobiSolver.TempScalarBuffer2,
                MyKernel.temperature_diffuse_weight_2D, MyKernel.temperature_diffuse_vector_2D,
                None, None)
            # GP.TemperatureInitialDiffuseWeights2D, GP.TemperatureInitialDiffuseVectors2D)

            DensityBuffer.swap()
            VelocityBuffer.swap()
            TemperatureBuffer.swap()
    elif SolverMode == SolverType.RBGS:
        RBGSSolver.GeneralRBGSSolver(RBGSSolver.Diffuse, VelocityBuffer)
        RBGSSolver.GeneralRBGSSolver(RBGSSolver.ScalarDiffuse, DensityBuffer, GP.DensityDiffuseFactor)
        RBGSSolver.GeneralRBGSSolver(RBGSSolver.ScalarDiffuse, TemperatureBuffer, GP.TemperatureDiffuseFactor)
    # w3 -> w4
    Math.Divergence(VelocityBuffer.cur, VelocityDivBuffer)
    if SolverMode == SolverType.Jacobi:
        # PressureBuffer.cur.fill(0)
        if GP.CompactPoisson == False:
            JacobiSolver.GeneralJacobiSolver(JacobiSolver.Poisson, PressureBuffer, VelocityDivBuffer)
        else:
            JacobiSolver.GeneralJacobiSolver(
                JacobiSolver.CompactPoisson, None,
                VelocityDivBuffer, PressureBuffer.cur, PressureBuffer.nxt,
                JacobiSolver.TempScalarBuffer1, JacobiSolver.TempScalarBuffer2,
                MyKernel.velocity_poisson_weight_2D, MyKernel.velocity_poisson_vector_2D,
                MyKernel.velocity_initial_poisson_weight_2D, MyKernel.velocity_initial_poisson_vector_2D)
            PressureBuffer.swap()
    elif SolverMode == SolverType.RBGS:
        RBGSSolver.GeneralRBGSSolver(RBGSSolver.Poisson, PressureBuffer, VelocityDivBuffer)
    Projection(VelocityBuffer.cur, PressureBuffer.cur)

    # Scalar
    # Advect(VelocityBuffer.cur, DensityBuffer.cur, DensityBuffer.nxt)
    # Advect(VelocityBuffer.cur, TemperatureBuffer.cur, TemperatureBuffer.nxt)
    # DensityBuffer.swap()
    # TemperatureBuffer.swap()

    Render(RenderingBuffer, TemperatureBuffer.cur, DensityBuffer.cur)


# Main Function
def Reset():
    VelocityBuffer.cur.fill(0)
    PressureBuffer.cur.fill(0)
    TemperatureBuffer.cur.fill(GP.AmbientTemperature)
    DensityBuffer.cur.fill(0)


def Main():
    # Local Class
    class VisualizeType(Enum):
        Dye = 0,
        Velocity = 1,
        Temperature = 2,

    # State Init
    Paused = False
    VisualizeMode = VisualizeType.Dye
    SolverMode = SolverType.Jacobi
    # Local Variable Init
    DisplayBuffer = Ti.Vector.field(3, float, shape=(GP.WindowResolution, GP.WindowResolution))

    # System Init
    Window = Ti.GUI('Stable Fluid', (GP.WindowResolution, GP.WindowResolution))
    MouseDataGen = MouseDataGenerator()

    # Ti.profiler.print_kernel_profiler_info('trace')
    # Ti.profiler.clear_kernel_profiler_info()  # clear all records
    Counter = 100

    # Main Loop
    while Window.running:
        # Handle Input
        if Window.get_event(Ti.GUI.PRESS):
            InputEvent = Window.event
            if InputEvent.key == Ti.GUI.ESCAPE:
                break
            elif InputEvent.key == 'r':
                Paused = False
                Reset()
            elif InputEvent.key == 's':
                if curl_strength:
                    curl_strength = 0
                else:
                    curl_strength = 7
            elif InputEvent.key == 'v':
                VisualizeMode = VisualizeType.Velocity
            elif InputEvent.key == 'd':
                VisualizeMode = VisualizeType.Dye
            elif InputEvent.key == 'c':
                visualize_c = True
                visualize_d = False
                visualize_v = False
            elif InputEvent.key == 'p':
                Paused = not Paused
            elif InputEvent.key == '1':
                SolverMode = SolverType.Jacobi
                print("SolverMode Switch To SolverType.Jacobi")
            elif InputEvent.key == '2':
                SolverMode = SolverType.RBGS
                print("SolverMode Switch To SolverType.RBGS")

        if not Paused:
            SystemStep(SolverMode)

        # if VisualizeMode == VisualizeType.Dye:
        # 	RenderingBuffer = DyeBuffer.cur
        # elif VisualizeMode == VisualizeType.Velocity:
        # 	Math.MulAdd(VelocityBuffer.cur, RenderingBuffer , 0.01, 0.5)

        Math.UpSample(RenderingBuffer, DisplayBuffer, GP.FluidsScale)
        Window.set_image(DisplayBuffer)
        Window.show()

    # if Counter == 0:
    # 	Paused = not Paused
    # 	# Ti.profiler.print_kernel_profiler_info('trace')
    # 	Ti.profiler.print_scoped_profiler_info()
    # Counter = Counter - 1


if __name__ == '__main__':
    Main()